import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

device = "cuda"

tokenizer = AutoTokenizer.from_pretrained("/data1/glm-4-9b-chat", trust_remote_code=True)


model = AutoModelForCausalLM.from_pretrained(
    "/data1/glm-4-9b-chat",
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=True,
    trust_remote_code=True
).to(device).eval()
for param in model.base_model.parameters():
    param.requires_grad = False
print(model.transformer.embedding)
for l in model.transformer.encoder.layers:
    print(model.transformer.embedding)

gen_kwargs = {"max_length": 2500, "do_sample": True, "top_k": 1}
with torch.no_grad():
    while True:
        print("您有什么要说的：\n")
        query = input()
        if query == "q":
            break

        inputs = tokenizer.apply_chat_template([{"role": "user", "content": query}],
                                            add_generation_prompt=True,
                                            tokenize=True,
                                            return_tensors="pt",
                                            return_dict=True
                                            )

        inputs = inputs.to(device)
        outputs = model.generate(**inputs, **gen_kwargs)
        outputs = outputs[:, inputs['input_ids'].shape[1]:]
        print(tokenizer.decode(outputs[0], skip_special_tokens=True))
